# bvh_tree.py

"""
BVH (Bounding Volume Hierarchy) 树的定义与构建。
"""

import numpy as np
import time
# from .constants import EPSILON # EPSILON 目前不直接在此文件使用

class BVHNode:
    """
    BVH (Bounding Volume Hierarchy) 树的节点。
    """
    def __init__(self, 
                 bbox_min: np.ndarray | None = None,
                 bbox_max: np.ndarray | None = None,
                 left_child: 'BVHNode | None' = None,
                 right_child: 'BVHNode | None' = None,
                 triangles: list[tuple[np.ndarray, np.ndarray, np.ndarray]] | None = None
                ):
        self.bbox_min = bbox_min
        self.bbox_max = bbox_max
        self.left_child = left_child
        self.right_child = right_child
        self.triangles = triangles if triangles is not None else []
        self.is_leaf = (left_child is None and right_child is None)

    def __repr__(self):
        if self.is_leaf:
            return f"BVHLeaf(triangles={len(self.triangles)}, bbox_min={self.bbox_min}, bbox_max={self.bbox_max})"
        else:
            return f"BVHInternalNode(bbox_min={self.bbox_min}, bbox_max={self.bbox_max})"

# --- BVH 构建辅助函数 ---

def get_triangle_bbox(
    triangle_vertices: tuple[np.ndarray, np.ndarray, np.ndarray]
) -> tuple[np.ndarray, np.ndarray]:
    """计算单个三角形的轴对齐包围盒 (AABB)。"""
    v0, v1, v2 = triangle_vertices
    min_coord = np.minimum(np.minimum(v0, v1), v2)
    max_coord = np.maximum(np.maximum(v0, v1), v2)
    return min_coord, max_coord

def get_triangle_centroid(
    triangle_vertices: tuple[np.ndarray, np.ndarray, np.ndarray]
) -> np.ndarray:
    """计算单个三角形的重心。"""
    v0, v1, v2 = triangle_vertices
    return (v0 + v1 + v2) / 3.0

# --- BVH 构建核心函数 ---

def _build_bvh_recursive(
    triangles: list[tuple[np.ndarray, np.ndarray, np.ndarray]],
    max_tris_per_leaf: int
) -> BVHNode:
    """递归构建BVH树的辅助函数。"""
    num_triangles = len(triangles)
    if num_triangles == 0:
        # 返回一个bbox为None的节点，上层traverse_bvh会处理
        return BVHNode(triangles=[], bbox_min=None, bbox_max=None) 

    # 1. 计算当前三角形集合的包围盒和重心
    current_node_bbox_min, current_node_bbox_max = get_triangle_bbox(triangles[0])
    centroids = [get_triangle_centroid(triangles[0])] 

    for i in range(1, num_triangles):
        tri_bbox_min, tri_bbox_max = get_triangle_bbox(triangles[i])
        current_node_bbox_min = np.minimum(current_node_bbox_min, tri_bbox_min)
        current_node_bbox_max = np.maximum(current_node_bbox_max, tri_bbox_max)
        centroids.append(get_triangle_centroid(triangles[i]))

    node = BVHNode(bbox_min=current_node_bbox_min, bbox_max=current_node_bbox_max)

    # 3. 终止条件：创建叶子节点
    if num_triangles <= max_tris_per_leaf:
        node.triangles = triangles
        node.is_leaf = True # 构造函数会推断，但显式设置更清晰
        return node

    # --- 非叶子节点，需要分割 ---
    node.is_leaf = False # 构造函数会推断，但显式设置更清晰

    # 4. 选择分割轴
    bbox_span = current_node_bbox_max - current_node_bbox_min
    split_axis = np.argmax(bbox_span)
    
    # 5. 分割三角形 (基于重心的均分法)
    sortable_triangles = []
    for i in range(num_triangles):
        sortable_triangles.append((centroids[i][split_axis], triangles[i]))
    sortable_triangles.sort(key=lambda x: x[0])
    
    sorted_triangles = [item[1] for item in sortable_triangles]
    mid_point_idx = num_triangles // 2
    left_triangles = sorted_triangles[:mid_point_idx]
    right_triangles = sorted_triangles[mid_point_idx:]

    if not left_triangles or not right_triangles:
        # print(f"警告 (BVH 构建): 分割失败于轴 {split_axis}。当前节点将成为叶子节点，含 {num_triangles} 三角形。")
        node.triangles = triangles 
        node.is_leaf = True
        return node

    # 6. 递归构建子节点
    node.left_child = _build_bvh_recursive(left_triangles, max_tris_per_leaf)
    node.right_child = _build_bvh_recursive(right_triangles, max_tris_per_leaf)
    node.triangles = [] # 内部节点不直接存储三角形
    return node

def build_bvh(
    all_model_triangles: list[tuple[np.ndarray, np.ndarray, np.ndarray]], 
    max_tris_per_leaf: int = 8,
    verbose: bool = False # verbose控制是否打印统计信息
) -> BVHNode | None:
    """构建BVH树的入口函数。"""
    if not all_model_triangles:
        # print("警告 (build_bvh): 输入的三角形列表为空，无法构建BVH。") # 日志由调用者处理
        return None
    
    root_node = _build_bvh_recursive(list(all_model_triangles), max_tris_per_leaf) # 使用副本

    if verbose:
        print_bvh_stats(root_node, model_total_triangles=len(all_model_triangles))
    return root_node

def print_bvh_stats(node: BVHNode | None, depth: int = 0, stats: dict | None = None, model_total_triangles: int = 0):
    """递归打印BVH树的一些统计信息 (用于调试和验证)"""
    if node is None:
        return

    if stats is None: # 首次调用时初始化计数器字典
        stats = {
            "node_count": 0,
            "leaf_count": 0,
            "total_tris_in_leaves": 0,
            "max_depth": -1 # 初始为-1，这样第一个深度0会成为max_depth
        }

    stats["node_count"] += 1
    stats["max_depth"] = max(stats["max_depth"], depth)

    if node.is_leaf:
        stats["leaf_count"] += 1
        tris_count = len(node.triangles)
        stats["total_tris_in_leaves"] += tris_count
    else:
        # 确保子节点存在才递归，尽管_build_bvh_recursive不应产生bbox为None的内部节点
        if node.left_child:
            print_bvh_stats(node.left_child, depth + 1, stats, model_total_triangles)
        if node.right_child:
            print_bvh_stats(node.right_child, depth + 1, stats, model_total_triangles)
    
    if depth == 0: # 如果是根节点的调用结束，打印最终统计
        print("\n--- BVH 树统计信息 ---")
        print(f"  总节点数量: {stats['node_count']}")
        print(f"  叶子节点数量: {stats['leaf_count']}")
        print(f"  最大深度: {stats['max_depth']}")
        if stats['leaf_count'] > 0:
             print(f"  平均每个叶子节点的三角形数量: {stats['total_tris_in_leaves'] / stats['leaf_count']:.2f}")
        else:
             print(f"  平均每个叶子节点的三角形数量: N/A (无叶子节点)")
        print(f"  叶子节点中总三角形数量: {stats['total_tris_in_leaves']}")
        if model_total_triangles > 0: # 只有当提供了模型总数时才打印对比
             print(f"  (模型总三角形数: {model_total_triangles})")
        print("-----------------------\n")


if __name__ == "__main__":
    print("--- 开始测试 bvh_tree.py ---")
    # 为了测试BVH构建，我们需要一个triangulate_face函数。
    # 理想情况下，这个测试应该独立，或者从voxelizer_core导入。
    # 这里我们临时定义一个简化的triangulate_face，仅用于此测试。
    # 或者，更好的做法是，如果voxelizer_core.py在同一包内，可以尝试导入。
    # 当使用 python -m objTojson_ts.bvh_tree 运行时，__package__ 会被设置
    # 此时相对导入应该可以工作
    try:
        from .voxelizer_core import triangulate_face # 使用相对导入
        print("  (已从 .voxelizer_core 成功导入 triangulate_face 进行测试)")
    except ImportError as e:
        # 如果上面的导入失败（例如直接 python bvh_tree.py），则打印更明确的错误
        print(f"  错误: 无法通过相对导入 '.voxelizer_core' 获取 triangulate_face。错误: {e}")
        print("        请尝试使用 'python -m objTojson_ts.bvh_tree' 从项目根目录的上一级运行测试。")
        print("        将跳过依赖 triangulate_face 的BVH测试部分。")
        # 如果导入失败，后续依赖 triangulate_face 的测试代码可能需要跳过或使用存根
        # 为了简单，如果导入失败，我们可能无法完整运行此处的BVH测试
        # model_triangles_for_bvh_build = [] # 确保后续代码不会因triangulate_face未定义而出错
        # 或者直接 exit()
        exit("Triangulate_face 导入失败，无法进行BVH构建测试。")
        def triangulate_face(face_indices: list[int], all_vertices: list[list[float]]
                             ) -> list[tuple[np.ndarray, np.ndarray, np.ndarray]]:
            # 这是一个非常简化的版本，仅用于测试数据生成
            triangles_np = []
            if len(face_indices) >= 3:
                v0 = np.array(all_vertices[face_indices[0]], dtype=np.float64)
                for i in range(1, len(face_indices) - 1):
                    v1 = np.array(all_vertices[face_indices[i]], dtype=np.float64)
                    v2 = np.array(all_vertices[face_indices[i+1]], dtype=np.float64)
                    triangles_np.append((v0, v1, v2))
            return triangles_np

    cube_vertices_list_for_bvh = [
        [0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0],
        [0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [1.0, 1.0, 1.0], [0.0, 1.0, 1.0]
    ]
    cube_faces_indices_list_for_bvh = [
        [0, 3, 2, 1], [4, 5, 6, 7], [0, 1, 5, 4], 
        [1, 2, 6, 5], [2, 3, 7, 6], [3, 0, 4, 7]
    ]

    model_triangles_for_bvh_build: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = []
    for face_idx_list in cube_faces_indices_list_for_bvh:
        triangles_from_face = triangulate_face(face_idx_list, cube_vertices_list_for_bvh)
        model_triangles_for_bvh_build.extend(triangles_from_face)
    
    print(f"用于BVH构建的测试模型共有 {len(model_triangles_for_bvh_build)} 个三角形。")

    if model_triangles_for_bvh_build:
        test_passed_count = 0
        test_total_count = 0
        for max_leaf_tris in [1, 2, 4, 8, 12, 13]: # 增加测试用例
            test_total_count +=1
            print(f"\n--- 测试构建BVH，max_tris_per_leaf = {max_leaf_tris} ---")
            bvh_start_time = time.time()
            bvh_root = build_bvh(model_triangles_for_bvh_build, 
                                 max_tris_per_leaf=max_leaf_tris, 
                                 verbose=False) # 不启用verbose打印统计
            bvh_end_time = time.time()
            print(f"  BVH构建耗时 (max_tris_per_leaf={max_leaf_tris}): {bvh_end_time - bvh_start_time:.4f} 秒。")
            
            current_test_passed = True
            if bvh_root:
                if not (bvh_root.bbox_min is not None and bvh_root.bbox_max is not None):
                    print(f"  错误: BVH根节点包围盒未正确设置 (max_tris_per_leaf={max_leaf_tris})")
                    current_test_passed = False
                else:
                    expected_root_min = np.array([0.0, 0.0, 0.0])
                    expected_root_max = np.array([1.0, 1.0, 1.0])
                    if not np.allclose(bvh_root.bbox_min, expected_root_min):
                        print(f"  错误: 根节点bbox_min {bvh_root.bbox_min} 与预期 {expected_root_min}不符 (max_tris_per_leaf={max_leaf_tris})")
                        current_test_passed = False
                    if not np.allclose(bvh_root.bbox_max, expected_root_max):
                        print(f"  错误: 根节点bbox_max {bvh_root.bbox_max} 与预期 {expected_root_max}不符 (max_tris_per_leaf={max_leaf_tris})")
                        current_test_passed = False
                
                # 验证叶子节点中的总三角形数量
                stats_check = {"total_tris_in_leaves": 0, "leaf_count": 0, "node_count": 0, "max_depth": -1}
                print_bvh_stats(bvh_root, stats=stats_check) # 调用一次以填充stats_check
                if stats_check["total_tris_in_leaves"] != len(model_triangles_for_bvh_build):
                    print(f"  错误: 叶子节点中总三角形数量 ({stats_check['total_tris_in_leaves']}) "
                          f"与模型总数 ({len(model_triangles_for_bvh_build)}) 不符 (max_tris_per_leaf={max_leaf_tris})")
                    current_test_passed = False

            else:
                print(f"  错误: BVH构建返回None (max_tris_per_leaf={max_leaf_tris})")
                current_test_passed = False
            
            if current_test_passed:
                print(f"  BVH (max_tris_per_leaf={max_leaf_tris}) 构建基本检查通过。")
                test_passed_count +=1
            else:
                print(f"  !!! BVH (max_tris_per_leaf={max_leaf_tris}) 构建基本检查失败 !!!")

        print(f"\n--- BVH 构建所有测试完成: {test_passed_count}/{test_total_count} 通过 ---")
    else:
        print("没有三角形用于BVH构建测试。")
    
    print("--- bvh_tree.py 测试结束 ---")